from flask import Flask, render_template, redirect, url_for, flash
from utils import allowed_doc_file, extract_doc_content
import os
import logging
from werkzeug.utils import secure_filename
from embedding_utils import get_doc_embedding
from milvus_utils import insert_doc_vector, search_doc_vectors

logger = logging.getLogger(__name__)

UPLOAD_FOLDER = os.getenv("UPLOAD_FOLDER", "uploads")
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)


# 处理文件上传的函数，参数请求对象request
def handle_doc_upload(request):
    if request.method == "POST":
        # 如果请求中没有文件
        if "file" not in request.files:
            # 则重定向回上一个页面
            return redirect(request.url)
        # 获取上传的文件对象
        file = request.files["file"]
        if file.filename == "":
            logger.warning("上传的文档为空")
            flash("上传的文档为空", "error")
            return redirect(url_for("upload_doc"))
        # 判断文件类型是否正确
        if file and allowed_doc_file(file.filename):
            try:
                file_name = secure_filename(file.filename)
                file_path = os.path.join(UPLOAD_FOLDER, file_name)
                file.save(file_path)
                logger.warning(f"文件已经成功保存到:{file_path}")
                # 闪现提示文件成功 把此消息写入会话
                flash(f"上传成功:{file.filename}", "success")
                # 提取文档的内容
                doc_content = extract_doc_content(file_path)
                logger.warning(f"文件内容提供成功，内容:{doc_content}")
                doc_embedding = get_doc_embedding(doc_content)
                logger.warning(f"向量计算成功,长度{len(doc_embedding)}")
                # insert_doc_vectors([doc_embedding], [doc_content])
                insert_doc_vector(doc_embedding, doc_content)
                return redirect(url_for("upload_doc"))
            except Exception as e:
                flash(f"上传失败，请重试:{file.filename}", "error")
                logger.error(f"文件保存失败:{e}")

    return render_template("upload_doc.html")


def handle_doc_search(request):
    results = []
    query = ""
    if request.method == "POST":
        query = request.form.get("query", "")
        logger.info(f"开始检索{query}")
        if query:
            # 获取查询的关键字对应的向量
            query_embedding = get_doc_embedding(query)
            # 在milvus里检索最相似的5个文档
            milvus_results = search_doc_vectors(query_embedding, limit=5)
            logger.warning(f"检索完成，命中的组数是{len(milvus_results)}")
            for hits in milvus_results:
                for hit in hits:
                    results.append(
                        {"doc": hit.entity.get("doc", ""), "score": hit.distance}
                    )
    return render_template("search_doc.html", results=results)
